import os
from mistralai import Mistral

# --- 配置 ---
# 使用您提供的固定API Key
API_KEY = "Kz2yZgyEQxkreQRagET0rtz15l8smbRT"
# 使用您指定的 Agent ID
AGENT_ID = "ag_019a3501501b768a901bef17061c2a20"

# --- 主程序 ---
try:
    # 1. 初始化客户端
    client = Mistral(api_key=API_KEY)
    print(f"客户端初始化成功，使用 Agent ID: {AGENT_ID}")
except Exception as e:
    print(f"初始化客户端时出错: {e}")
    exit()

# 2. 初始化一个列表来存储对话历史
conversation_history = []

print("\n多轮对话机器人已启动 (Agent Stream 模式)。")
print("输入 'exit' 或 'quit' 退出。")
print("---")

while True:
    try:
        # 3. 获取用户输入
        user_input = input("你: ")
        if user_input.lower() in ["exit", "quit"]:
            print("再见！")
            break
        
        if not user_input.strip(): # 如果用户只输入了空格或回车，则跳过
            continue

        # 4. 将用户输入添加到对话历史中
        conversation_history.append({"role": "user", "content": user_input})

        # 5. 发送完整的对话历史到 API
        response_stream = client.beta.conversations.start_stream(
            inputs=conversation_history,
            agent_id=AGENT_ID,
        )
        
        # 6. 处理数据流并拼接完整回复
        assistant_full_response = ""
        print("模型: ", end="", flush=True) # 准备在同一行打印流式输出

        for chunk in response_stream:
            # 我们只关心包含消息内容增量的事件
            if chunk.event == 'message.output.delta':
                # 实时打印数据块内容，实现打字机效果
                print(chunk.data.content, end="", flush=True)
                # 将数据块内容拼接到完整回复字符串中
                assistant_full_response += chunk.data.content
        
        print() # 流式输出结束后换行

        # 7. 将模型的完整回复也添加到对话历史中，为下一轮对话做准备
        if assistant_full_response:
             conversation_history.append({"role": "assistant", "content": assistant_full_response})

    except Exception as e:
        # 8. 捕获并打印错误信息
        print(f"\n[错误] 程序运行或API调用时出错: {e}")
        
        # 如果出错，移除刚刚添加的用户消息，以便用户可以重新尝试发送
        if conversation_history and conversation_history[-1]["role"] == "user":
            conversation_history.pop()
        
        print("---")

